# Copyright (c) 2010
# The Regents of the University of Michigan
# All Rights Reserved

# Permission is granted to use, copy, create derivative works, and
# redistribute this software and such derivative works for any purpose,
# so long as the name of the University of Michigan is not used in
# any advertising or publicity pertaining to the use or distribution
# of this software without specific, written prior authorization. If
# the above copyright notice or any other identification of the
# University of Michigan is included in any copy of any portion of
# this software, then the disclaimer below must also be included.

# This software is provided as is, without representation or warranty
# of any kind either express or implied, including without limitation
# the implied warranties of merchantability, fitness for a particular
# purpose, or noninfringement.  The Regents of the University of
# Michigan shall not be liable for any damages, including special,
# indirect, incidental, or consequential damages, with respect to any
# claim arising out of or in connection with the use of the software,
# even if it has been or is hereafter advised of the possibility of
# such damages.

'''lookup.py -- simple Kademlia lookup implementation.

Supports standalone use of PyDHT for playing wth Vuze.
'''

import bisect
import heapq
import itertools
import logging
import socket
import sys

import event

import cvuze_util
from ordered_set import OrderedListSet
from pydht import contact, dht, dhtutil, routing_table


LOOKUP_CONCURRENCY = 5  # The alpha parameter from the Kademlia paper.
LOOKUP_K = 20  # The k parameter from the Kademlia paper.
MAX_RETRIES = 1


BOOTSTRAP_NODES = [(socket.gethostbyname('dht.vuze.com'), 6881)]


class PendingFindNode(object):
  '''Bookkeeping information for a pending FIND-NODE request in a Lookup.
  '''

  def __init__(self, parent, contact):
    self.parent = parent
    self.contact = contact
    self.tries_left = MAX_RETRIES
    self.timeout = event.timeout(10, self._query_timeout)

  def _query_timeout(self):
    lookup = self.parent
    client = lookup.parent

    if lookup.lookup_id is None:
      return

    if self.tries_left > 0:
      # Retry.
      self.tries_left -= 1

      client.send_find_node(self.contact, lookup.lookup_id)
      event.timeout(10, self._query_timeout)
    else:
      self.parent.router.didnt_respond(self.contact)
      # Remove this pending find and fire more.
      try:
        self.parent.pending_finds.remove(self)
      except ValueError:
        pass
      for _ in xrange(len(lookup.pending_finds), LOOKUP_CONCURRENCY):
        if not lookup._send_query():
          break


class Lookup(object):
  '''An ongoing "lookup", in the sense defined by the Kademlia paper.
  '''
  
  def __init__(self, parent, lookup_id, query_fn, finished_callback, address, router):
    self.parent = parent
    self.query_fn = query_fn
    self.lookup_id = lookup_id
    self.finished_callback = finished_callback
    self.address = address
    self.router = router

    def distance_from_target_id(contact):
      return cvuze_util.compute_id_distance(lookup_id,
                                            dht.compute_node_id(contact))

    # Set of IDs of contacts that have replied, ordered by distance
    # from this instance. A heap would do, but we'll use an
    # OrderedListSet because we need that for contacts_to_query anyway.
    self.replied_contacts = OrderedListSet([],
        lambda x: distance_from_target_id(x[0]))

    # Set of IDs of contacts we've queried.
    self.contacts_queried = set()

    # Set of contacts we still need to query. Must be kept sorted by
    # contact ID's distance from the target ID.
    self.contacts_to_query = OrderedListSet([], distance_from_target_id)

    self.pending_finds = []
    for node in BOOTSTRAP_NODES:
      self.router.record_live(node)  # TODO: decide whether router should hide Contact.
    self._start()

  def _start(self):
    '''Start the lookup.'''
    to_query = self.router.get_k_closest(self.lookup_id)
    logging.debug('starting with %d nodes' % len(to_query))
    for _, node in itertools.izip(
        xrange(min(LOOKUP_CONCURRENCY, len(to_query))),
        itertools.cycle(to_query)):
      self.contacts_to_query.add(node)

    logging.debug('starting with %d contacts' % len(self.contacts_to_query))
    for _ in xrange(LOOKUP_CONCURRENCY):
      if not self._send_query():
        break

  def _should_continue(self):
    if self.lookup_id is None:
      return False

    if not self.replied_contacts:
      return True

    if not self.contacts_to_query:
      self._finish()
      return False

    # Lookup halts when we have no one left to query that's closer
    # than the farthest responder.
    closest_dist = cvuze_util.compute_id_distance(
        dht.compute_node_id(self.contacts_to_query[0]), self.lookup_id)
    farthest_responder_dist = cvuze_util.compute_id_distance(
        dht.compute_node_id(self.replied_contacts[-1][0]), self.lookup_id)
    
    if (closest_dist > farthest_responder_dist and
        len(self.replied_contacts) >= LOOKUP_K):
      self._finish()
      return False
    else:
      return True

  def _send_query(self):
    '''Start another FIND-NODE query.
    If there are >= LOOKUP_CONCURRENCY queries running already, no-op.
    Also no-op if there are no contacts to query right now.

    Return True if we started another query, False otherwise.
    '''

    logging.debug('send query called with %d finds, %d to query' % (len(self.pending_finds), len(self.contacts_to_query)))
    # Too many pending FIND-NODEs already?
    if len(self.pending_finds) == LOOKUP_CONCURRENCY:
      return False

    # Lookup finished?
    if self.lookup_id is None:
      return False

    # No one to query?
    if not self.contacts_to_query:
      return False

    # Should finish?
    if not self._should_continue():
      return False

    to_query = self.contacts_to_query[0]
    del self.contacts_to_query[0]
    logging.debug('Querying node with ID %s' % dht.compute_hex_node_id(*to_query))

    self.query_fn(to_query, self.lookup_id)
    self.contacts_queried.add(dht.compute_node_id(to_query))
    self.pending_finds.append(PendingFindNode(self, to_query))

    return True

  def handle_find_node_reply(self, reply, contact_address):
    if self.lookup_id is None:
      return False

    replying_contact = contact_address

    logging.debug('Reply from node with ID %r' % dht.compute_node_id(replying_contact))
    # Mark the replier as such.
    self.replied_contacts.add((replying_contact, reply.spoof_id))

    # We only need the K closest replying contacts.
    if len(self.replied_contacts) > LOOKUP_K:
      del self.replied_contacts[-1]

    for idx, elt in enumerate(self.pending_finds):
      if elt.contact == replying_contact:
        del self.pending_finds[idx]
        break

    for con in reply.contacts:
      if con == self.address:
        continue
      if (dht.compute_node_id(con) in self.contacts_queried or
          con in self.contacts_to_query):
        continue
      self.contacts_to_query.add(con)
      self.router.record_live(con, live=False)

    if self._should_continue():
      for _ in xrange(len(self.pending_finds), LOOKUP_CONCURRENCY):
        if not self._send_query():
          break

  def _finish(self):
    self.lookup_id = None
    self.finished_callback(self.replied_contacts.elts)


class LookupValue(Lookup):
  def __init__(self, *args):
    self.values = []
    self.waiting_for_continuation = True
    Lookup.__init__(self, *args)

  def handle_find_value_reply(self, reply, contact_address):
    if reply.values:
      logging.debug('got reply with values from %s' % dht.compute_hex_node_id(*contact_address))
      self.values.extend(reply.values)
      if not reply.has_continuation:
        self.waiting_for_continuation = False
        self._finish()
    else:
      logging.debug('%s did not have it' % dht.compute_hex_node_id(*contact_address))
      self.handle_find_node_reply(# XXX: DHT-dependent
                                  dht.serialization.ReplyFindNode(reply.header, 0, 0, 0,
                                                reply.contacts),
                                  contact_address)

  def _finish(self):
    self.lookup_id = None
    self.finished_callback(self.values)

  def _should_continue(self):
    if self.lookup_id is None:
      return False
    if not self.waiting_for_continuation:
      self._finish()
      return False
    return Lookup._should_continue(self)


class LookupClient(routing_table.RouterAwareClient):
  '''A DHT client whose FIND-NODE requests are for lookups only, not probing.

  This is the "normal" case of using a DHT, rather than doing research on it.
  '''

  def __init__(self, address, router):
    dht.DHTClient.__init__(self, address)
    self.the_lookup = None
    self.set_router(router)

  def lookup(self, id_to_find, is_value=False):
    '''Synchronously perform a Kademlia lookup for the given ID.

    Not re-entrant.
    '''
    # Synchrony avoids, for now, the problem of multiple concurrent lookups.
    # After it works, I'll make it re-entrant.
    out_nodes = []
    def lookup_done(nodes):
      out_nodes[:] = nodes
      event.abort()

    if is_value:
      query_fn = self.send_find_value
      lookup_class = LookupValue
    else:
      query_fn = self.send_find_node
      lookup_class = Lookup
    self.the_lookup = lookup_class(self, id_to_find, query_fn, lookup_done,
                                   self.address, self.router)
    event.dispatch()
    self.the_lookup = None
    return out_nodes

  def lookup_value(self, id_to_find):
    return self.lookup(id_to_find, True)

  # sender_contact is the same as sender_address; some other DHTs have
  # more contact information.
  def handle_find_node_reply(self, reply, sender_contact, sender_address):
    if self.the_lookup is None:
      if self.mbs:
        for con in reply.contacts:
          self.router.record_live(con, False)
      return
    self.the_lookup.handle_find_node_reply(reply, sender_contact)

  def handle_find_value_reply(self, reply, sender_contact, sender_address):
    if self.the_lookup is not None:
      self.the_lookup.handle_find_value_reply(reply, sender_contact)

  def minibootstrap(self):
    self.mbs = True
    self.send_find_node(BOOTSTRAP_NODES[0], self.router.local_contact.dht_id)
    event.timeout(1, event.abort)
    event.dispatch()
    self.mbs = False


if __name__ == '__main__':
  import datetime
  me = (sys.argv[1], 1234)
  c = LookupClient(me, routing_table.RoutingTable(me))

  print 'Starts at %s' % datetime.datetime.today()
  print c.router.get_k_closest(dht.compute_node_id(me))
  res = c.lookup(dht.compute_node_id(me))
  print 'First finished at %s' % datetime.datetime.today()
  print 'lookup result: %s' % str(sorted([dht.compute_hex_node_id(x[0]) for x in res]))
  print 'closest nodes: %s' % str(sorted([dht.compute_hex_node_id(x) for x in c.router.get_k_closest(dht.compute_node_id(me))]))
  res = c.lookup(dht.compute_node_id(me))
  print 'Second finished at %s' % datetime.datetime.today()
  print 'lookup result: %s' % str(sorted([dht.compute_hex_node_id(x[0]) for x in res]))
  print 'closest nodes: %s' % str(sorted([dht.compute_hex_node_id(x) for x in c.router.get_k_closest(vuze.compute_node_id(me))]))

